// Package checksum checksum/calculator.go
package checksum

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"hash"
	"hash/crc32"
	"px-table-checksum-go/db" // 引入db包以使用ChecksumTask和FormatForCRC
	"px-table-checksum-go/models"
	"px-table-checksum-go/util" // 导入 util 包以使用 FormatForCRC
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
)

// 定义错误常量，便于统一错误处理和测试
var (
	// ErrNilDB 数据库连接为空错误
	ErrNilDB = errors.New("数据库连接为空")

	// ErrNilTask 任务为空错误
	ErrNilTask = errors.New("校验任务为空")

	// ErrInvalidKeyInfo 键信息无效错误
	ErrInvalidKeyInfo = errors.New("无效的键信息")

	// ErrQueryTimeout 查询超时错误
	ErrQueryTimeout = errors.New("查询执行超时")
)

// 对象池，提高性能
var (
	// IEEE CRC32多项式表 (使用IEEE标准，与MySQL CRC32()函数兼容)
	crc32Table = crc32.MakeTable(crc32.IEEE)

	// CRC32哈希计算器池，复用哈希对象减少内存分配
	crcHasherPool = sync.Pool{
		New: func() interface{} {
			return crc32.New(crc32Table)
		},
	}

	// 字符串构建器池，复用构建器减少内存分配
	stringBuilderPool = sync.Pool{
		New: func() interface{} {
			return new(strings.Builder)
		},
	}
)

// CalculateChunkChecksum 计算给定数据块的校验和
//
// 该函数执行数据块的校验计算，获取指定数据块的记录，计算它们的校验和。
// 它返回包含校验结果的 ChecksumResult，一个表示是否是最后一个块的标志，以及可能的错误。
//
// 参数:
//   - dbConn: 用于执行查询的数据库连接 (*sql.DB)
//   - task: 校验任务，包含表名、键信息和查询参数
//
// 返回值:
//   - *models.ChecksumResult: 校验结果，包含CRC32校验值、行数等信息
//   - bool: 是否是最后一个块
//   - error: 计算过程中的错误
func CalculateChunkChecksum(dbConn *sql.DB, task *db.ChecksumTask) (*models.ChecksumResult, bool, error) {
	// 参数验证
	if dbConn == nil {
		return nil, false, ErrNilDB
	}
	if task == nil {
		return nil, false, ErrNilTask
	}
	// 校验任务中的关键信息
	if task.KeyInfo == nil || len(task.KeyInfo.Columns) == 0 {
		return nil, false, fmt.Errorf("%w: 缺少键列定义 (表: %s.%s)", ErrInvalidKeyInfo, task.SchemaName, task.TableName)
	}
	if task.SQLTemplate == "" {
		return nil, false, errors.New("查询模板为空")
	}

	// 创建带上下文信息的日志记录器
	logger := log.WithFields(log.Fields{
		"schema":  task.SchemaName,
		"table":   task.TableName,
		"chunk":   task.ChunkNo,
		"db_id":   task.DBID,
		"min_key": task.StartKeyString,
	})

	logger.Debug("开始计算数据块校验和")

	// 准备查询参数
	queryParams, err := prepareQueryParams(task)
	if err != nil {
		return nil, false, fmt.Errorf("构建查询参数失败: %w", err)
	}

	// 获取查询超时时间 (默认30秒, 可以从配置读取)
	// TODO: 从配置中获取超时时间 config.GlobalConfig.Options.QueryTimeout (假设有这个配置项)
	timeoutSeconds := 30

	// 创建带超时的上下文
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeoutSeconds)*time.Second)
	defer cancel() // 确保上下文最终被取消，释放资源

	// 执行查询
	logger.WithFields(log.Fields{
		"sql":  task.SQLTemplate,
		"args": fmt.Sprintf("%v", queryParams), // 记录查询参数
	}).Debug("执行校验查询")

	rows, err := dbConn.QueryContext(ctx, task.SQLTemplate, queryParams...)
	if err != nil {
		// 检查是否是超时错误
		if ctx.Err() == context.DeadlineExceeded {
			logger.Errorf("查询执行超时 (%ds)", timeoutSeconds)
			return nil, false, fmt.Errorf("%w: %v", ErrQueryTimeout, err)
		}
		logger.Errorf("执行校验查询失败: %v", err)
		return nil, false, fmt.Errorf("执行校验查询失败: %w", err)
	}
	// 确保查询结果集被关闭
	defer rows.Close()

	// 处理查询结果
	result, isLastChunk, err := processQueryRows(ctx, rows, task)
	if err != nil {
		// processQueryRows内部已经记录了错误日志
		return nil, false, fmt.Errorf("处理查询结果失败: %w", err)
	}

	// 记录执行结果
	if result != nil {
		logger.WithFields(log.Fields{
			"rows":    result.RowCount,
			"crc32":   result.CRC32,
			"max_key": result.MaxKey,
			"is_last": isLastChunk,
		}).Info("数据块校验计算完成")
	} else {
		// 这种情况理论上不应该发生，因为 processQueryRows 会返回一个空结果对象
		logger.Warn("数据块无内容 (或处理结果为nil)，校验计算完成")
	}

	return result, isLastChunk, nil
}

// prepareQueryParams 准备查询参数
// 根据任务的 StartKeyValues 或 StartKeyString 生成用于 SQL 查询的参数切片
func prepareQueryParams(task *db.ChecksumTask) ([]interface{}, error) {
	// 如果有解析好的 StartKeyValues，直接使用
	if task.StartKeyValues != nil && len(task.StartKeyValues) > 0 {
		// 验证键值数量是否与键列数匹配
		if len(task.StartKeyValues) != task.NumKeyColumns {
			return nil, fmt.Errorf("键值数 %d 与键列数 %d 不匹配",
				len(task.StartKeyValues), task.NumKeyColumns)
		}
		return task.StartKeyValues, nil
	}

	// 如果没有 StartKeyValues，尝试从 StartKeyString 解析
	// 检查 StartKeyString 是否为空，这通常表示第一个块
	if task.StartKeyString == "" || task.StartKeyString == util.GetInitialStartKey(task.NumKeyColumns) {
		// 对于第一个块，使用 GetInitialStartKey 对应的解析结果
		// Modify: Use util.ParseKeyFromString and util.GetInitialStartKey
		return util.ParseKeyFromString(util.GetInitialStartKey(task.NumKeyColumns), task.NumKeyColumns)
	}

	// 解析 StartKeyString 为参数
	// *** 修改点: 使用 util 包中的 ParseKeyFromString ***
	values, err := util.ParseKeyFromString(task.StartKeyString, task.NumKeyColumns)
	if err != nil {
		return nil, fmt.Errorf("解析起始键值 '%s' 失败: %w", task.StartKeyString, err)
	}
	// Modify: Return values directly, no need for db.GetSQLArgumentValuesFromKeys unless specific driver needs
	return values, nil
}

// processQueryRows 处理查询结果集
// 遍历数据库查询返回的行，计算CRC32校验和，并确定是否为最后一个块。
func processQueryRows(ctx context.Context, rows *sql.Rows, task *db.ChecksumTask) (*models.ChecksumResult, bool, error) {
	logger := log.WithFields(log.Fields{
		"schema":  task.SchemaName,
		"table":   task.TableName,
		"chunk":   task.ChunkNo,
		"db_id":   task.DBID,
		"min_key": task.StartKeyString,
	})

	// 获取列名，用于后续处理
	columns, err := rows.Columns()
	if err != nil {
		logger.Errorf("获取列名失败: %v", err)
		return nil, false, fmt.Errorf("获取列名失败: %w", err)
	}
	if len(columns) == 0 {
		logger.Error("查询结果没有列名")
		return nil, false, errors.New("查询结果没有列名")
	}

	// 准备列类型信息，用于创建正确的扫描目标
	colTypes, err := rows.ColumnTypes()
	if err != nil {
		logger.Errorf("获取列类型信息失败: %v", err)
		return nil, false, fmt.Errorf("获取列类型信息失败: %w", err)
	}

	// 跟踪行统计和CRC32计算
	var (
		rowCount     int
		maxKeyValues []interface{} // 记录块内最后一行的键值
		crc32Value   uint32        // 最终的CRC32校验和
	)

	// 从配置或常量获取期望的块大小 (用于判断是否是最后一个块)
	// TODO: 从配置读取 config.GlobalConfig.Options.ChunkSize
	expectedChunkSize := 1000 // 默认值

	// 遍历结果行
	for rows.Next() {
		// 定期检查上下文是否已取消（允许长时间查询的提前取消）
		if rowCount > 0 && rowCount%100 == 0 { // 每处理100行检查一次
			select {
			case <-ctx.Done():
				logger.Warnf("处理查询结果时操作被取消: %v", ctx.Err())
				return nil, false, fmt.Errorf("处理查询结果时操作被取消: %w", ctx.Err())
			default:
				// 继续处理
			}
		}

		// 创建接收查询结果的变量 (使用列类型信息)
		values := make([]interface{}, len(columns)) // 存储转换后的 Go 值
		values = values
		valuePtrs := make([]interface{}, len(columns)) // 存储指向扫描目标的指针
		for i, colType := range colTypes {
			// *** 修改点: 使用 db 包中的 CreateScanTarget ***
			scanTarget := db.CreateScanTarget(colType.DatabaseTypeName()) // 获取扫描目标指针
			if scanTarget == nil {
				// 理论上不应该发生，除非 CreateScanTarget 未处理某种类型
				errMsg := fmt.Sprintf("无法为列 %s (类型 %s) 创建扫描目标", columns[i], colType.DatabaseTypeName())
				logger.Error(errMsg)
				return nil, false, errors.New(errMsg)
			}
			valuePtrs[i] = scanTarget
		}

		// 扫描行数据到指针指向的目标
		if err := rows.Scan(valuePtrs...); err != nil {
			logger.Errorf("扫描行数据失败 (第 %d 行): %v", rowCount+1, err)
			return nil, false, fmt.Errorf("扫描行数据失败 (第 %d 行): %w", rowCount+1, err)
		}

		// 从扫描目标指针获取实际的 Go 值
		scannedValues := make([]interface{}, len(columns)) // 存储转换后的值
		for i := 0; i < len(columns); i++ {
			// *** 修改点: 使用 db 包中的 GetScannedValue ***
			val, scanErr := db.GetScannedValue(valuePtrs[i]) // 获取转换后的值
			if scanErr != nil {
				// 处理获取值时的错误
				errMsg := fmt.Sprintf("获取列 %s 的扫描值失败: %v", columns[i], scanErr)
				logger.Error(errMsg)
				return nil, false, errors.New(errMsg)
			}
			scannedValues[i] = val // 存储转换后的值
		}

		// 更新行计数
		rowCount++

		// 计算此行数据的CRC32值并累加
		// 使用优化版本，利用对象池减少内存分配
		rowCRC32 := calculateRowCRC32Optimized(scannedValues) // 使用扫描到的 Go 值
		// CRC32校验和的合并通常使用异或(XOR)操作
		// 注意：简单的异或可能不是最佳的合并方式，取决于具体的校验需求。
		// 如果需要严格的块校验和，应该在 calculateRowCRC32Optimized 内部完成整个块的计算。
		// 但如果目标是每行CRC异或，则保持如下：
		crc32Value = crc32Value ^ rowCRC32 // 使用异或合并 (确认这是否是期望的合并逻辑)
		// 如果目标是 Percona Toolkit 的行为，它可能不是简单异或，而是将行数据流式写入 hasher。
		// 为简单起见，这里暂时保留异或，但实际应用中需要确认算法。

		// 更新最大键值（记录当前行的键值，循环结束后即为最后一行的键值）
		// 确保只取键列的值
		numKeyCols := task.NumKeyColumns
		if numKeyCols <= 0 || numKeyCols > len(scannedValues) {
			logger.Errorf("无效的键列数配置: %d (总列数: %d)", numKeyCols, len(scannedValues))
			return nil, false, fmt.Errorf("无效的键列数配置: %d", numKeyCols)
		}
		maxKeyValues = make([]interface{}, numKeyCols)
		for i := 0; i < numKeyCols; i++ {
			// 键列在 SELECT 语句中被放在了前面，索引 i 对应键列
			// Modify: Access key values using their position in the SELECT statement
			maxKeyValues[i] = scannedValues[i] // 从扫描到的 Go 值中获取键值
		}
	} // 结束 rows.Next() 循环

	// 检查行遍历过程中的错误
	if err := rows.Err(); err != nil {
		logger.Errorf("遍历行数据时发生错误: %v", err)
		return nil, false, fmt.Errorf("遍历行数据时发生错误: %w", err)
	}

	// 检查是否为最后一个块（查询无行 或 行数少于期望的块大小）
	// 注意: 这里的判断逻辑可能需要根据实际业务调整。
	// 例如，即使行数少于 expectedChunkSize，如果查询没有返回错误且有数据，
	// 并且知道后面没有更多数据了（例如通过主键范围判断），才应标记为 isLastChunk。
	// 当前的简单判断：无行或者行数不足预期则认为是最后一个块。
	isLastChunk := rowCount == 0 || rowCount < expectedChunkSize

	// 当没有行时，返回一个表示空块的结果
	if rowCount == 0 {
		logger.Info("查询结果为空，返回空块结果")
		return &models.ChecksumResult{
			DBID:       task.DBID,
			SchemaName: task.SchemaName,
			TableName:  task.TableName,
			ChunkNo:    task.ChunkNo,
			CreateTime: time.Now(),
			MinKey:     task.StartKeyString,
			MaxKey:     task.StartKeyString, // 对于空结果，最大键等于最小键
			CRC32:      0,                   // 空块CRC32为0
			RowCount:   0,
		}, true, nil // rowCount 为 0 时，isLastChunk 应为 true
	}

	// 格式化最大键值为字符串，用于存储
	// *** 修改点: 使用 util 包中的 FormatKeyToString ***
	maxKeyString, err := util.FormatKeyToString(maxKeyValues)
	if err != nil {
		logger.Errorf("格式化最大键值失败: %v, 键值: %v", err, maxKeyValues)
		// 如果格式化失败，可能无法正确计算下一个块的起始键，应返回错误
		return nil, false, fmt.Errorf("格式化最大键值失败: %w", err)
	}

	// 创建结果对象
	result := &models.ChecksumResult{
		DBID:       task.DBID,
		SchemaName: task.SchemaName,
		TableName:  task.TableName,
		ChunkNo:    task.ChunkNo,
		CreateTime: time.Now(), // 记录结果生成时间
		MinKey:     task.StartKeyString,
		MaxKey:     maxKeyString,
		CRC32:      crc32Value,
		RowCount:   rowCount,
	}

	return result, isLastChunk, nil
}

// calculateRowCRC32Optimized 使用对象池优化的行CRC32计算
// 将一行所有列的值格式化为字节，并计算CRC32校验和。
func calculateRowCRC32Optimized(rowValues []interface{}) uint32 {
	// 验证输入参数
	if rowValues == nil {
		log.Warn("计算CRC32时收到空行值，返回0")
		return 0
	}

	// 从池中获取字符串构建器
	builder := stringBuilderPool.Get().(*strings.Builder)
	// 确保构建器最终被重置并放回池中
	defer func() {
		builder.Reset()
		stringBuilderPool.Put(builder)
	}()

	// 从池中获取CRC32哈希计算器
	hasher := crcHasherPool.Get().(hash.Hash32)
	// 确保哈希器最终被重置并放回池中
	defer func() {
		hasher.Reset()
		crcHasherPool.Put(hasher)
	}()

	// 遍历行中的每个值
	for i, val := range rowValues {
		// 在值之间添加分隔符 (例如 '#')，确保不同列的值不会粘连
		// Modify: Use a constant separator defined in util or checksum package
		if i > 0 {
			builder.WriteByte('#') // 使用 '#' 作为分隔符，可配置
		}
		// 使用 util.FormatForCRC 将 Go 值转换为字节切片
		bytesToWrite := util.FormatForCRC(val) // 调用移至 util 的函数
		builder.Write(bytesToWrite)            // 将格式化后的字节写入构建器
	}

	// 获取构建好的完整行字符串的字节表示
	rowBytes := []byte(builder.String())

	// 将行字节写入哈希计算器
	// hash.Hash 的 Write 方法保证不返回错误
	_, _ = hasher.Write(rowBytes)

	// 计算并返回最终的32位CRC校验和
	return hasher.Sum32()
}
